#!/usr/bin/env python3
"""
初始化示例数据
添加YOLO11s模型和示例服务器
"""

import asyncio
import time
import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from src.storage import data_manager, ModelInfo, SubscriptionServer
from src.utils import generate_id
from src.logger import logger


async def init_sample_data():
    """初始化示例数据"""
    logger.info("Initializing sample data...")
    
    current_time = time.time()
    
    # 1. 添加YOLO11s模型
    yolo11s_model = ModelInfo(
        id=generate_id(),
        name="yolo11s",
        version="1.0.0",
        model_type="detection",
        description="YOLO11s - 小型目标检测模型，平衡速度和精度",
        author="Ultralytics",
        size="s",
        upload_time=current_time,
        last_update=current_time,
        status="available",
        download_url="https://github.com/ultralytics/assets/releases/download/v8.2.0/yolo11s.pt",
        checksum="a1b2c3d4e5f6",
        server_id="local",
        tags=["detection", "yolo", "small", "fast"],
        requirements={
            "python": ">=3.8",
            "torch": ">=1.8.0",
            "torchvision": ">=0.9.0",
            "ultralytics": ">=8.0.0"
        },
        metrics={
            "mAP50": 0.537,
            "mAP50-95": 0.374,
            "speed_cpu": "1.2ms",
            "speed_gpu": "0.8ms",
            "params": "11.1M",
            "flops": "28.4G"
        }
    )
    
    # 2. 添加YOLO11n模型（更小的版本）
    yolo11n_model = ModelInfo(
        id=generate_id(),
        name="yolo11n",
        version="1.0.0",
        model_type="detection",
        description="YOLO11n - 纳米级目标检测模型，极致轻量化",
        author="Ultralytics",
        size="xs",
        upload_time=current_time,
        last_update=current_time,
        status="available",
        download_url="https://github.com/ultralytics/assets/releases/download/v8.2.0/yolo11n.pt",
        checksum="b2c3d4e5f6a1",
        server_id="local",
        tags=["detection", "yolo", "nano", "lightweight"],
        requirements={
            "python": ">=3.8",
            "torch": ">=1.8.0",
            "torchvision": ">=0.9.0",
            "ultralytics": ">=8.0.0"
        },
        metrics={
            "mAP50": 0.496,
            "mAP50-95": 0.336,
            "speed_cpu": "0.8ms",
            "speed_gpu": "0.5ms",
            "params": "2.6M",
            "flops": "6.5G"
        }
    )
    
    # 3. 添加YOLO11m模型（中等版本）
    yolo11m_model = ModelInfo(
        id=generate_id(),
        name="yolo11m",
        version="1.0.0",
        model_type="detection",
        description="YOLO11m - 中型目标检测模型，精度和速度的最佳平衡",
        author="Ultralytics",
        size="m",
        upload_time=current_time,
        last_update=current_time,
        status="available",
        download_url="https://github.com/ultralytics/assets/releases/download/v8.2.0/yolo11m.pt",
        checksum="c3d4e5f6a1b2",
        server_id="local",
        tags=["detection", "yolo", "medium", "balanced"],
        requirements={
            "python": ">=3.8",
            "torch": ">=1.8.0",
            "torchvision": ">=0.9.0",
            "ultralytics": ">=8.0.0"
        },
        metrics={
            "mAP50": 0.591,
            "mAP50-95": 0.431,
            "speed_cpu": "2.4ms",
            "speed_gpu": "1.5ms",
            "params": "20.1M",
            "flops": "68.2G"
        }
    )
    
    # 4. 添加示例服务器
    sample_server = SubscriptionServer(
        id=generate_id(),
        name="本地YOLO服务器",
        url="http://localhost:7880",
        api_key=None,
        status="online",
        last_check=current_time,
        check_interval=1800,
        consecutive_errors=0
    )
    
    try:
        # 检查模型是否已存在
        existing_models = await data_manager.get_models()
        existing_model_names = {model.name for model in existing_models}
        
        # 添加模型（如果不存在）
        models_to_add = [yolo11s_model, yolo11n_model, yolo11m_model]
        added_models = 0
        
        for model in models_to_add:
            if model.name not in existing_model_names:
                await data_manager.add_model(model)
                logger.info(f"Added model: {model.name}")
                added_models += 1
            else:
                logger.info(f"Model {model.name} already exists, skipping")
        
        # 检查服务器是否已存在
        existing_servers = await data_manager.get_servers()
        existing_server_names = {server.name for server in existing_servers}
        
        # 添加服务器（如果不存在）
        if sample_server.name not in existing_server_names:
            await data_manager.add_server(sample_server)
            logger.info(f"Added server: {sample_server.name}")
        else:
            logger.info(f"Server {sample_server.name} already exists, skipping")
        
        logger.info(f"Sample data initialization completed. Added {added_models} models.")
        
        # 显示当前数据统计
        all_models = await data_manager.get_models()
        all_servers = await data_manager.get_servers()
        
        print(f"\n=== 数据统计 ===")
        print(f"模型总数: {len(all_models)}")
        print(f"服务器总数: {len(all_servers)}")
        
        print(f"\n=== 可用模型 ===")
        for model in all_models:
            print(f"- {model.name} v{model.version} ({model.model_type}, {model.size})")
            print(f"  描述: {model.description}")
            print(f"  mAP50: {model.metrics.get('mAP50', 'N/A')}")
            print()
        
        print(f"=== 服务器列表 ===")
        for server in all_servers:
            print(f"- {server.name}: {server.url} (状态: {server.status})")
        
    except Exception as e:
        logger.error(f"Failed to initialize sample data: {e}")
        raise


async def clear_sample_data():
    """清除示例数据"""
    logger.info("Clearing sample data...")
    
    try:
        # 获取所有数据
        models = await data_manager.get_models()
        servers = await data_manager.get_servers()
        
        # 删除YOLO模型
        yolo_models = [m for m in models if m.name.startswith('yolo11')]
        for model in yolo_models:
            await data_manager.delete_model(model.id)
            logger.info(f"Deleted model: {model.name}")
        
        # 删除示例服务器
        sample_servers = [s for s in servers if s.name == "本地YOLO服务器"]
        for server in sample_servers:
            await data_manager.delete_server(server.id)
            logger.info(f"Deleted server: {server.name}")
        
        logger.info("Sample data cleared successfully")
        
    except Exception as e:
        logger.error(f"Failed to clear sample data: {e}")
        raise


async def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description="Initialize or clear sample data")
    parser.add_argument(
        "action",
        choices=["init", "clear", "status"],
        help="Action to perform"
    )
    
    args = parser.parse_args()
    
    if args.action == "init":
        await init_sample_data()
    elif args.action == "clear":
        await clear_sample_data()
    elif args.action == "status":
        # 显示当前数据状态
        models = await data_manager.get_models()
        servers = await data_manager.get_servers()
        
        print(f"当前数据状态:")
        print(f"- 模型数量: {len(models)}")
        print(f"- 服务器数量: {len(servers)}")
        
        if models:
            print(f"\n模型列表:")
            for model in models:
                print(f"  - {model.name} v{model.version}")
        
        if servers:
            print(f"\n服务器列表:")
            for server in servers:
                print(f"  - {server.name}: {server.url}")


if __name__ == "__main__":
    asyncio.run(main())
